package com.explment.lens.utils.rsa;

import javax.crypto.Cipher;
import java.io.*;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.security.*;
import java.security.spec.*;
import java.util.Base64;

public class RSAUtils {
    private static final String ALGORITHM = "RSA";
    private static final String SIGN_ALGORITHM = "SHA256withRSA";
    private static final int KEY_SIZE = 2048;

    private PrivateKey privateKey;
    private PublicKey publicKey;

    // 从文件加载密钥对
    public void loadKeyPair(String publicKeyPath, String privateKeyPath) throws Exception {
        this.publicKey = (PublicKey) readKeyFromFile(publicKeyPath, true);
        this.privateKey = (PrivateKey) readKeyFromFile(privateKeyPath, false);
    }

    // 生成并保存新的密钥对到文件
    public void generateAndSaveKeyPair(String publicKeyPath, String privateKeyPath) throws Exception {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(ALGORITHM);
        keyPairGenerator.initialize(KEY_SIZE);
        KeyPair keyPair = keyPairGenerator.generateKeyPair();

        this.publicKey = keyPair.getPublic();
        this.privateKey = keyPair.getPrivate();

        saveKeyToFile(publicKeyPath, this.publicKey, true);
        saveKeyToFile(privateKeyPath, this.privateKey, false);
    }

    // 私钥签名
    public byte[] sign(String data) throws Exception {
        if (privateKey == null) {
            throw new IllegalStateException("Private key is not initialized");
        }

        Signature signature = Signature.getInstance(SIGN_ALGORITHM);
        signature.initSign(privateKey);
        signature.update(data.getBytes(StandardCharsets.UTF_8));
        return signature.sign();
    }

    // 公钥验签
    public boolean verify(String data, String sign) throws Exception {
        if (publicKey == null) {
            throw new IllegalStateException("Public key is not initialized");
        }

        Signature signature = Signature.getInstance(SIGN_ALGORITHM);
        signature.initVerify(publicKey);
        signature.update(data.getBytes(StandardCharsets.UTF_8));
        return signature.verify(Base64.getDecoder().decode(sign));
    }

    // 公钥加密
    public byte[] encrypt(String data) throws Exception {
        if (publicKey == null) {
            throw new IllegalStateException("Public key is not initialized");
        }

        Cipher cipher = Cipher.getInstance(ALGORITHM);
        cipher.init(Cipher.ENCRYPT_MODE, publicKey);
        return cipher.doFinal(data.getBytes(StandardCharsets.UTF_8));
    }

    // 私钥解密
    public String decrypt(byte[] encryptedData) throws Exception {
        if (privateKey == null) {
            throw new IllegalStateException("Private key is not initialized");
        }

        Cipher cipher = Cipher.getInstance(ALGORITHM);
        cipher.init(Cipher.DECRYPT_MODE, privateKey);
        return new String(cipher.doFinal(encryptedData), StandardCharsets.UTF_8);
    }

    // 读取密钥文件
    private Key readKeyFromFile(String filePath, boolean isPublic) throws Exception {
        byte[] keyBytes = Files.readAllBytes(new File(filePath).toPath());
        KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM);

        if (isPublic) {
            X509EncodedKeySpec keySpec = new X509EncodedKeySpec(keyBytes);
            return keyFactory.generatePublic(keySpec);
        } else {
            PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(keyBytes);
            return keyFactory.generatePrivate(keySpec);
        }
    }

    // 保存密钥到文件
    private void saveKeyToFile(String filePath, Key key, boolean isPublic) throws Exception {
        byte[] keyBytes = key.getEncoded();
        Files.write(new File(filePath).toPath(), keyBytes);
    }

    // 使用公钥加密字节数组（用于加密AES密钥）
    public byte[] encryptWithPublicKey(byte[] data) throws Exception {
        if (publicKey == null) {
            throw new IllegalStateException("Public key is not initialized");
        }
        Cipher cipher = Cipher.getInstance(ALGORITHM);
        cipher.init(Cipher.ENCRYPT_MODE, publicKey);
        return cipher.doFinal(data);
    }

    // 使用私钥解密字节数组（用于解密AES密钥）
    public byte[] decryptWithPrivateKey(byte[] encryptedData) throws Exception {
        if (privateKey == null) {
            throw new IllegalStateException("Private key is not initialized");
        }
        Cipher cipher = Cipher.getInstance(ALGORITHM);
        cipher.init(Cipher.DECRYPT_MODE, privateKey);
        return cipher.doFinal(encryptedData);
    }


    /*public static void main(String[] args) throws Exception {
        RSAUtils rsaUtils = new RSAUtils();

        try {
            // 1. 生成并保存密钥对
            System.out.println("=== 生成密钥对 ===");
            rsaUtils.generateAndSaveKeyPair(
                    "src/main/resources/keys/public.key",
                    "src/main/resources/keys/private.key"
            );

            // 2. 加载密钥对
            System.out.println("\n=== 加载密钥对 ===");
            rsaUtils.loadKeyPair(
                    "src/main/resources/keys/public.key",
                    "src/main/resources/keys/private.key"
            );

            // 3. 加密解密测试
            System.out.println("\n=== 加密解密测试 ===");
            String secretMessage = "这是高度机密的信息123!@#";
            System.out.println("原始消息: " + secretMessage);

            byte[] encrypted = rsaUtils.encrypt(secretMessage);
            System.out.println("加密后(Base64): " + Base64.getEncoder().encodeToString(encrypted));

            String decrypted = rsaUtils.decrypt(encrypted);
            System.out.println("解密后: " + decrypted);

            // 4. 签名验签测试
            System.out.println("\n=== 签名验签测试 ===");
            String dataToSign = "重要交易数据: 金额=1000元, 时间=2023-01-01";
            System.out.println("待签名数据: " + dataToSign);

            byte[] sig = rsaUtils.sign(dataToSign);
            System.out.println("签名(Base64): " + Base64.getEncoder().encodeToString(sig));

            boolean valid = rsaUtils.verify(dataToSign, Base64.getEncoder().encodeToString(sig));
            System.out.println("签名验证: " + (valid ? "成功" : "失败"));

        } catch (Exception e) {
            System.err.println("处理过程中出错: " + e.getMessage());
            e.printStackTrace();
        }
    }*/

}

